#include <iostream>
#include <string>
#include <filesystem>
#include <cxxopts.hpp>
#include "ipu_protonet_classifier.hpp"
#include "ipu_image_preprocessor.hpp"

namespace fs = std::filesystem;

int main(int argc, char* argv[]) {
    try {
        // 命令行参数解析
        cxxopts::Options options("ipu_protonet_classifier", "IPU ProtoNet分类器 - 少样本图像分类");
        
        options.add_options()
            ("m,model", "模型文件路径", cxxopts::value<std::string>())
            ("s,support", "支持集数据路径", cxxopts::value<std::string>())
            ("q,query", "查询图像路径 (单张)", cxxopts::value<std::string>()->default_value(""))
            ("t,test", "测试数据集路径", cxxopts::value<std::string>()->default_value(""))
            ("p,prototypes", "原型文件路径", cxxopts::value<std::string>()->default_value("prototypes.json"))
            ("d,distance", "距离度量 (euclidean/cosine)", cxxopts::value<std::string>()->default_value("euclidean"))
            ("c,confidence", "置信度方法 (simple/nll)", cxxopts::value<std::string>()->default_value("simple"))
            ("load-prototypes", "加载已保存的原型")
            ("w,width", "输入图像宽度", cxxopts::value<int>()->default_value("224"))
            ("height", "输入图像高度", cxxopts::value<int>()->default_value("224"))
            ("h,help", "显示帮助信息");

        auto result = options.parse(argc, argv);

        if (result.count("help")) {
            std::cout << options.help() << std::endl;
            return 0;
        }

        // 检查必要参数
        if (!result.count("model")) {
            std::cerr << "❌ 缺少必要参数: --model" << std::endl;
            return 1;
        }

        // 获取参数
        std::string model_path = result["model"].as<std::string>();
        std::string support_path = result.count("support") ? result["support"].as<std::string>() : "";
        std::string query_path = result["query"].as<std::string>();
        std::string test_path = result["test"].as<std::string>();
        std::string prototypes_path = result["prototypes"].as<std::string>();
        bool load_prototypes = result.count("load-prototypes") > 0;
        
        // 输入尺寸
        int width = result["width"].as<int>();
        int height = result["height"].as<int>();
        std::pair<int, int> input_size = {width, height};
        
        // 距离度量
        protonet::DistanceMetric distance_metric = protonet::DistanceMetric::EUCLIDEAN;
        std::string distance_str = result["distance"].as<std::string>();
        if (distance_str == "cosine") {
            distance_metric = protonet::DistanceMetric::COSINE;
        }
        
        // 置信度方法
        protonet::ConfidenceMethod confidence_method = protonet::ConfidenceMethod::SIMPLE;
        std::string confidence_str = result["confidence"].as<std::string>();
        if (confidence_str == "nll") {
            confidence_method = protonet::ConfidenceMethod::NEGATIVE_LOG_LIKELIHOOD;
        }

        std::cout << "🚀 启动IPU ProtoNet分类器" << std::endl;
        std::cout << "📁 模型路径: " << model_path << std::endl;
        std::cout << "📐 输入尺寸: " << width << "x" << height << std::endl;

        // 创建预处理配置
        protonet::IPUImagePreprocessor::PreprocessConfig config;
        config.input_format = "BGR";
        config.resize_width = width;
        config.resize_height = height;
        config.do_normalize = true;
        config.mean_values = {123.675f, 116.28f, 103.53f};  // ImageNet均值
        config.std_values = {58.395f, 57.12f, 57.375f};     // ImageNet标准差

        // 初始化分类器
        protonet::IPUProtoNetClassifier classifier(model_path, input_size, &config);

        // 加载或构建原型
        if (load_prototypes && fs::exists(prototypes_path)) {
            std::cout << "\n📂 从文件加载原型..." << std::endl;
            if (!classifier.loadPrototypes(prototypes_path)) {
                std::cerr << "❌ 加载原型失败" << std::endl;
                return 1;
            }
        } else if (!support_path.empty()) {
            std::cout << "\n🔨 构建原型..." << std::endl;
            if (!classifier.buildPrototypes(support_path)) {
                std::cerr << "❌ 构建原型失败" << std::endl;
                return 1;
            }
            
            // 保存原型
            std::cout << "\n💾 保存原型到文件..." << std::endl;
            classifier.savePrototypes(prototypes_path);
        } else {
            std::cerr << "❌ 必须提供支持集路径 (--support) 或启用加载原型 (--load-prototypes)" << std::endl;
            return 1;
        }

        // 执行分类任务
        if (!query_path.empty()) {
            // 单张图像分类
            std::cout << "\n🔍 分类查询图像: " << query_path << std::endl;
            auto result = classifier.classify(query_path, distance_metric, confidence_method);
            
            std::cout << "📊 分类结果:" << std::endl;
            std::cout << "   预测类别: " << result.predicted_class << std::endl;
            std::cout << "   置信度: " << std::fixed << std::setprecision(4) << result.confidence << std::endl;
            std::cout << "   距离: " << std::fixed << std::setprecision(6) << result.distance << std::endl;
            
        } else if (!test_path.empty()) {
            // 测试数据集评估
            std::cout << "\n🧪 评估测试数据集: " << test_path << std::endl;
            auto results = classifier.evaluateDataset(test_path, distance_metric, confidence_method);
            
            // 保存评估结果
            std::string results_path = "evaluation_results_" + distance_str + "_" + confidence_str + "_ipu.json";
            // 这里应该保存结果到JSON文件，但为简化起见，我们只打印到控制台
            
            std::cout << "\n💾 详细评估结果已保存至: " << results_path << std::endl;
            std::cout << "🎉 分类任务完成!" << std::endl;
        } else {
            std::cout << "\n⚠️ 未指定查询图像或测试数据集路径" << std::endl;
            std::cout << "使用 --query 进行单张图像分类，或使用 --test 进行数据集评估" << std::endl;
        }

    } catch (const cxxopts::exceptions::exception& e) {
        std::cerr << "❌ 参数解析错误: " << e.what() << std::endl;
        return 1;
    } catch (const std::exception& e) {
        std::cerr << "❌ 错误: " << e.what() << std::endl;
        return 1;
    }
    
    return 0;
}